import cv2
import numpy as np
import json
import os


def generate_labelme_json(segmented_image_path, output_json_path, rgb_colors, image_path='', file_name='image.jpg'):
    """
    将分割图像转换为 labelme 格式的标注文件

    :param segmented_image_path: 分割图像的路径
    :param output_json_path: 要保存的 JSON 文件路径
    :param rgb_colors: 字典类型，建立色彩值和标签名称之间的映射关系
    :param image_path: 图像路径
    :param file_name: 图像文件名
    """
    img = cv2.imread(segmented_image_path)
    label_data = {'version': '4.5.7', 'flags': {}, 'shapes': [], 'imagePath': file_name, 'imageData': None}
    label_data['imageHeight'], label_data['imageWidth'], label_data['imageChannels'] = img.shape
    # if image_path:
    #     label_data['imagePath'] = file_name

    for color, label in rgb_colors.items():
        mask = cv2.inRange(img, np.array(color), np.array(color))
        cv2.imwrite("mask.png",mask)
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            polygon_vertices = contour.reshape(-1, 2).tolist()
            shape_data = {'label': label, 'points': polygon_vertices, 'group_id': None, 'shape_type': 'polygon',
                          'flags': {}}
            label_data['shapes'].append(shape_data)

    with open(output_json_path, 'w') as file:
        json.dump(label_data, file, indent=4)


# 测试
rgb_colors = {
    (0, 0, 128): '_background_',
    (0, 128, 0): 'leaf',
    (0, 128, 128): 'harmArea'
}
labelDirPath = r"F:\tmp\paddle\data\230823\result\pseudo_color_prediction"
dirPath = r"F:\tmp\paddle\data\230823\result\pseudo_color_prediction"
imageNames = os.listdir(labelDirPath)
for imageName in imageNames:
    suffix = imageName.split(".")[-1]
    name_list = imageName.split(".")
    name_list.pop()
    name_no_suffix = "".join(name_list)
    segmented_image_path = os.path.join(labelDirPath,imageName)
    output_json_path = os.path.join(dirPath,name_no_suffix+".json")
    image_path = os.path.join(dirPath,name_no_suffix+".jpg")
    generate_labelme_json(segmented_image_path, output_json_path, rgb_colors, image_path=image_path,
                          file_name=os.path.basename(image_path))

print("完成")